package gold.digger;

import gold.utils.InputUtil;

import java.util.*;

/**
 * Created by fanzhenyu02 on 2021/12/10.
 * common problem solver template.
 */
public class LC2049 {
    public long startExecuteTime = System.currentTimeMillis();


    class Solution {

        public class TreeNode {
            public int val;
            public TreeNode father;
            public TreeNode left;
            public TreeNode right;
            public int weight = 0;

            public TreeNode(int x) {
                val = x;
            }

            TreeNode() {
            }
        }

        public int countHighestScoreNodes(int[] parents) {
            int n = parents.length;
            TreeNode[] fullTree = new TreeNode[n];
            TreeNode root = null;

            for (int i = 0; i < n; i++) {
                if (fullTree[i] == null) fullTree[i] = new TreeNode(i);
                if (parents[i] == -1) root = fullTree[i];
                else {
                    if (fullTree[parents[i]] == null) fullTree[parents[i]] = new TreeNode(parents[i]);
                    if (null == fullTree[parents[i]].left) fullTree[parents[i]].left = fullTree[i];
                    else fullTree[parents[i]].right = fullTree[i];
                    fullTree[i].father = fullTree[parents[i]];
                }
            }

            dfsGetWeight(root);
            TreeMap<Long, Integer> cnt = new TreeMap<>(Comparator.reverseOrder());
            dfsGetPlusWeight(root, root.weight, cnt);
            return cnt.values().iterator().next();
        }

        public int dfsGetWeight(TreeNode root) {
            if (null == root) return 0;
            root.weight = dfsGetWeight(root.left) + dfsGetWeight(root.right);
            return root.weight + 1;
        }

        public void dfsGetPlusWeight(TreeNode root, int topWeight, TreeMap<Long, Integer> cnt) {
            if (null == root) return;
            int upperPart = null == root.father ? 1 : (topWeight - root.weight);
            int leftPart = null == root.left ? 1 : (root.left.weight + 1);
            int rightPart = null == root.right ? 1 : (root.right.weight + 1);
            long curRes = 1L * upperPart * leftPart * rightPart;
            cnt.put(curRes, cnt.getOrDefault(curRes, 0) + 1);
            dfsGetPlusWeight(root.left, topWeight, cnt);
            dfsGetPlusWeight(root.right, topWeight, cnt);
        }
    }

    class Solution_Other {

        long maxScore = 0;
        int ans = 0;
        int n;

        public int countHighestScoreNodes(int[] parents) {
            // 删除一个节点最多把整颗树分割成三部分：左子树、右子树、父节点及父节点的另一半子树
            // 所以，我们可以遍历每个节点的左右子树的数目，父节点及父节点的另一半子树的数量就等于 总节点数 减去 左右子树的数目 再减 一
            // 三者相乘就是分数，没有的部分用 1 代替
            // 而我们需要先构造出来这颗树才能通过DFS遍历
            this.n = parents.length;

            TreeNode[] nodes = new TreeNode[n];
            for (int i = 0; i < n; i++) {
                nodes[i] = new TreeNode();
            }

            for (int i = 1; i < n; i++) {
                nodes[parents[i]].addChild(nodes[i]);
            }

            dfs(nodes[0]);

            return ans;
        }

        private int dfs(TreeNode node) {
            if (node == null) {
                return 0;
            }

            int leftCount = dfs(node.left);
            int rightCount = dfs(node.right);
            int remain = n - leftCount - rightCount - 1;

            long score = help(leftCount) * help(rightCount) * help(remain);
            if (score == maxScore) {
                ans++;
            } else if (score > maxScore) {
                maxScore = score;
                ans = 1;
            }

            return leftCount + rightCount + 1;
        }

        private long help(int count) {
            return count == 0 ? 1 : count;
        }

        class TreeNode {
            TreeNode left;
            TreeNode right;

            void addChild(TreeNode child) {
                if (this.left == null) {
                    this.left = child;
                } else {
                    this.right = child;
                }
            }
        }
    }

    public void run() {
//        System.out.println(new Solution().countHighestScoreNodes(InputUtil.toIntegerArray("[-1,2,0,2,0]")));
//        System.out.println(new Solution().countHighestScoreNodes(InputUtil.toIntegerArray("[-1,2,0]")));
        System.out.println(new Solution().countHighestScoreNodes(InputUtil.toIntegerArray("[-1,98509,71017,36635,66652,32285,35669,71437,27157,24152,21448,1034,73576,3640,81534,94480,4909,3588,74245,63302,12796,25232,24916,6284]")));
    }

    public static void main(String[] args) throws Exception {
        LC2049 an = new LC2049();
        an.run();

        System.out.println("\ncurrent solution total execute time: " + (System.currentTimeMillis() - an.startExecuteTime) + " ms.");
    }
}
